import tensorflow as tf
from tensorflow.keras import layers, models
import numpy as np
import os
from PIL import Image
from typing import Tuple, List
import logging

logging.basicConfig(level=logging.INFO)

class CaptchaModelTrainer:
    def __init__(self, 
                 image_height: int = 64,
                 image_width: int = 128,
                 char_length: int = 4):
        """
        初始化验证码模型训练器
        Args:
            image_height: 图片高度
            image_width: 图片宽度
            char_length: 验证码长度
        """
        self.image_height = image_height
        self.image_width = image_width
        self.char_length = char_length
        self.characters = '0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
        self.num_classes = len(self.characters)
        self.model = self._build_model()

    def _build_model(self) -> models.Model:
        """
        构建CNN模型
        Returns:
            tf.keras.Model: 构建好的模型
        """
        input_layer = layers.Input(shape=(self.image_height, self.image_width, 3))
        
        # CNN layers
        x = layers.Conv2D(32, (3, 3), activation='relu', padding='same')(input_layer)
        x = layers.MaxPooling2D((2, 2))(x)
        x = layers.Conv2D(64, (3, 3), activation='relu', padding='same')(x)
        x = layers.MaxPooling2D((2, 2))(x)
        x = layers.Conv2D(128, (3, 3), activation='relu', padding='same')(x)
        x = layers.MaxPooling2D((2, 2))(x)
        
        # Dense layers
        x = layers.Flatten()(x)
        x = layers.Dense(512, activation='relu')(x)
        x = layers.Dropout(0.3)(x)
        
        # Output layers - one for each character position
        outputs = []
        for _ in range(self.char_length):
            output = layers.Dense(self.num_classes, activation='softmax')(x)
            outputs.append(output)
        
        model = models.Model(inputs=input_layer, outputs=outputs)
        model.compile(
            optimizer='adam',
            loss=['categorical_crossentropy'] * self.char_length,
            metrics=['accuracy']
        )
        
        return model

    def _preprocess_image(self, image_path: str) -> np.ndarray:
        """
        预处理图片
        Args:
            image_path: 图片路径
        Returns:
            np.ndarray: 预处理后的图片数组
        """
        image = Image.open(image_path)
        image = image.resize((self.image_width, self.image_height))
        image = np.array(image) / 255.0
        return image

    def _encode_label(self, text: str) -> List[np.ndarray]:
        """
        将文本标签转换为one-hot编码
        Args:
            text: 验证码文本
        Returns:
            List[np.ndarray]: one-hot编码后的标签列表
        """
        encoded = []
        for char in text:
            char_idx = self.characters.find(char)
            one_hot = np.zeros(self.num_classes)
            one_hot[char_idx] = 1
            encoded.append(one_hot)
        return encoded

    def prepare_dataset(self, 
                       data_dir: str) -> Tuple[np.ndarray, List[np.ndarray]]:
        """
        准备训练数据集
        Args:
            data_dir: 数据集目录
        Returns:
            Tuple[np.ndarray, List[np.ndarray]]: 处理后的图片和标签
        """
        images = []
        labels = [[] for _ in range(self.char_length)]
        
        for filename in os.listdir(data_dir):
            if not filename.endswith('.png'):
                continue
                
            # 从文件名获取标签
            label = filename.split('_')[0]
            if len(label) != self.char_length:
                continue
                
            image_path = os.path.join(data_dir, filename)
            image = self._preprocess_image(image_path)
            encoded_label = self._encode_label(label)
            
            images.append(image)
            for i, l in enumerate(encoded_label):
                labels[i].append(l)
        
        return np.array(images), [np.array(l) for l in labels]

    def train(self, 
             train_dir: str,
             validation_dir: str = None,
             epochs: int = 50,
             batch_size: int = 32) -> None:
        """
        训练模型
        Args:
            train_dir: 训练数据目录
            validation_dir: 验证数据目录
            epochs: 训练轮数
            batch_size: 批次大小
        """
        logging.info("准备训练数据...")
        x_train, y_train = self.prepare_dataset(train_dir)
        
        validation_data = None
        if validation_dir:
            logging.info("准备验证数据...")
            x_val, y_val = self.prepare_dataset(validation_dir)
            validation_data = (x_val, y_val)
        
        logging.info("开始训练模型...")
        self.model.fit(
            x_train, y_train,
            validation_data=validation_data,
            epochs=epochs,
            batch_size=batch_size,
            callbacks=[
                tf.keras.callbacks.EarlyStopping(
                    patience=5,
                    restore_best_weights=True
                ),
                tf.keras.callbacks.ModelCheckpoint(
                    'best_model.h5',
                    save_best_only=True
                )
            ]
        )
        
        logging.info("保存模型...")
        self.model.save('captcha_model.h5')

    def evaluate(self, test_dir: str) -> None:
        """
        评估模型
        Args:
            test_dir: 测试数据目录
        """
        logging.info("准备测试数据...")
        x_test, y_test = self.prepare_dataset(test_dir)
        
        logging.info("评估模型...")
        results = self.model.evaluate(x_test, y_test)
        
        for i, (loss, acc) in enumerate(zip(results[::2], results[1::2])):
            logging.info(f"Position {i+1} - Loss: {loss:.4f}, Accuracy: {acc:.4f}") 